package com.markus.code.并查集;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * Author:markusZhang
 * Date:Create in 2020/8/14 20:38
 * todo: 实现并查集的框架
 */
public class UnionAndSetDemo {
    public static class Element<V>{
        V value;
        public Element(V value){
            this.value = value;
        }
    }
    public static class UnionAndSet<V>{
        private HashMap<Element<V>,Element<V>> fatherMap;
        private HashMap<V,Element<V>> elementMap;
        private HashMap<Element<V>,Integer> sizeMap;
        public UnionAndSet(List<V> list){
            this.fatherMap = new HashMap<>();
            this.elementMap = new HashMap<>();
            this.sizeMap = new HashMap<>();
            //初始化集合，每个元素都是一个集合
            for (V v : list) {
                Element<V> element = new Element<>(v);
                elementMap.put(v,element);
                fatherMap.put(element,element);
                sizeMap.put(element,1);
            }
        }
        //得到两个元素所属集合的头结点
        private Element<V> getHeader(Element<V> element){
            Stack<Element<V>> stack = new Stack<>();
            Element<V> father = fatherMap.get(element);
            while(element != father){
                stack.push(element);//沿途所有节点入栈
                element = father;
                father = fatherMap.get(element);
            }
            //退出循环表示找到头结点了，这时候需要做的事情是将沿途所有节点都指向这个父节点
            while(!stack.isEmpty()){
                fatherMap.put(stack.pop(),father);
            }
            return father;
        }
        //合并两个节点
        public void union(V v1,V v2){
            //合并的前提需要两个元素都存在elementMap里
            if (elementMap.containsKey(v1) && elementMap.containsKey(v2)){
                Element<V> father_v1 = getHeader(elementMap.get(v1));
                Element<V> father_v2 = getHeader(elementMap.get(v2));
                if (father_v1 == father_v2){
                    return ;
                }
                //只有在头节点不同的情况下才进行连通
                Element<V> big = father_v1;
                Element<V> small = father_v2;
                int size_v1 = sizeMap.get(big);
                int size_v2 = sizeMap.get(small);
                big = size_v1 < size_v2 ? small : big;
                small = big==father_v1?father_v2:father_v1;
                int size = sizeMap.get(big)+sizeMap.get(small);
                fatherMap.put(small,big);
                sizeMap.put(big,size);
            }
        }
        public boolean isSameSet(V v1,V v2){
            if (elementMap.containsKey(v1) && elementMap.containsKey(v2)){
                Element<V> header_v1 = getHeader(elementMap.get(v1));
                Element<V> header_v2 = getHeader(elementMap.get(v2));
                return header_v1 == header_v2;
            }
            return false;
        }
    }

    public static void main(String[] args) {
        List<Integer> list = new ArrayList<>();
        list.add(1);
        list.add(2);
        list.add(3);
        list.add(4);
        UnionAndSet<Integer> unionAndSet = new UnionAndSet<>(list);
        unionAndSet.union(1,2);
        unionAndSet.union(3,4);
        System.out.println(unionAndSet.isSameSet(1,3));
        unionAndSet.union(2,3);
        System.out.println(unionAndSet.isSameSet(1,3));
    }
}
